#include "common.cuh"
#include "utilities.cuh"
#include "inversion.cuh"
#include "blocklu.cuh"
#include "lutrtri.cuh"
#include "cpu_blas.h"
#include "gpu_blas.cuh"
namespace CPU {
void update_blocklu(double* A, double* fractor, lapack_int * ipiv, int m, int n, int lda, int ldf) {
    // solve U12 through L11 @ U12 = P11.T @ A12 
    memcpy_submatrix(&A[m * lda], &fractor[m * ldf], m, n - m, lda, ldf);
    trsm('L', 'L', 'U', fractor, &fractor[m * ldf], m, n - m, ldf, ldf);

    // solve L21 through L21 @ U11 = A21, equivalent to U11.T @ L21.T = A21.T
    memcpy_submatrix(&A[m], &fractor[m], n - m, m, lda, ldf);
    trsm('R', 'U', 'N', fractor, &fractor[m], n - m, m, ldf, ldf);

    // compute \hat{A22} = A22 - L21 @ U12 (Schur Complement)
    memcpy_submatrix(&A[m + m * lda], &fractor[m + m * ldf], n - m, n - m, lda, ldf);
    gemm(&fractor[m], &fractor[m * ldf], &fractor[m + m * ldf], n - m, n - m, m, ldf, ldf, ldf, -1.0, 1.0);
    
    // LU decompose the h_A22
    getrf(&fractor[m + m * ldf], ipiv, n - m, n - m, ldf);
    for(int i = 0; i < n - m; i++) {
        assert(ipiv[i] == i + 1);
    }
}   

};

namespace DEVICE{

void update_blocklu(SOLVER_HANDLE handle, double* A, double* fractor, int m, int n, int lda, int ldf) {
    // solve U12 through L11 @ U12 = P11.T @ A12 
    memcpy_submatrix(&A[m * lda], &fractor[m * ldf], m, n - m, lda, ldf);
    trsm(handle.cublasH, 'L', 'L', 'U', fractor, &fractor[m * ldf], m, n - m, ldf, ldf);

    // solve L21 through L21 @ U11 = A21, equivalent to U11.T @ L21.T = A21.T
    memcpy_submatrix(&A[m], &fractor[m], n - m, m, lda, ldf);
    trsm(handle.cublasH, 'R', 'U', 'N', fractor, &fractor[m], n - m, m, ldf, ldf);

    // compute \hat{A22} = A22 - L21 @ U12 (Schur Complement)
    memcpy_submatrix(&A[m + m * lda], &fractor[m + m * ldf], n - m, n - m, lda, ldf);
    gemm(handle.cublasH, &fractor[m], &fractor[m * ldf], &fractor[m + m * ldf], n - m, n - m, m, ldf, ldf, ldf, -1.0, 1.0);
    
    // LU decompose the h_A22
    getrf(handle.cudenseH, handle.workspace, &fractor[m + m * ldf], n - m, n - m, ldf, handle.devInfo);
    CUDA_RT_CALL(cudaDeviceSynchronize());
}

};